from kuai.molio import read_mol, write_mol
from kuai.ffio import read_ff, write_ff

import os
import os.path

DFF_OPT_DFI = """$SYSTEM
   ROOT = %(DFF_ROOT)s
$END
$MODEL
   %(TEMP_FILE_NAME)s.msd
$END
$FORCEFIELD
   TYPE = AMBER
   INPUT = %(TEMP_FILE_NAME)s.ppf
   USE_CUTOFF = YES
   NBMETHOD = ATOM_BASED
   CUTOFF = 8.5
   BUFFER = 0.5
   TAIL_CORRECTION = YES
   EQUI = INPUT
$END
$MMOPTIONS
   CONVERGENCE = %(CONVERGENCE)e
   MAXSTEPS = 200
   METHOD = %(OPT_METHOD)s
$END
$PRINT
   CONSOLEXYZ = NO
$END
"""

def optimize(mol, ff, vars = {}):
    if "DFF_BIN" not in vars:
        vars["DFF_BIN"] = r"F:\Madlee\Job\Direct Force Field\7\bin32w"
    if "DFF_ROOT" not in vars:
        vars["DFF_ROOT"] = os.path.dirname(vars["DFF_BIN"])
    if "TEMP_FILE_NAME" not in vars:
        vars["TEMP_FILE_NAME"] = ".dff"
        
    if "OPT_METHOD" not in vars:
        vars["OPT_METHOD"] = "BFGS"
        
    if "CONVERGENCE" not in vars:
        vars["CONVERGENCE"] = 1.0e-8

    dfi = open(vars["TEMP_FILE_NAME"] + ".dfi", "w")
    dfi.write(DFF_OPT_DFI % vars)
    dfi.close()
    write_mol(vars["TEMP_FILE_NAME"] + ".msd", mol)
    # write_ff(vars["TEMP_FILE_NAME"] + ".ppf", ff)
    
    command = '"' + os.path.join(vars["DFF_BIN"], "dffmm") + '" '
    command +=  " " + vars["TEMP_FILE_NAME"]
    pin, pout, perr = os.popen3(command)
    pin.close()
    pout.read()
    error = perr.read().strip()
    if error != "":
        raise RuntimeError(error)
    else:
        return read_mol(vars["TEMP_FILE_NAME"] + ".msd")

